﻿#ifndef __WPE_TRUNK__
#define __WPE_TRUNK__

#include <map>
#include "thisinc.h"
#include "wpearch.h"
#include "strdump.h"
#include "libstr.hpp"
#include "bwm.h"
#include "crc64.h"

namespace wpe
{
	enum TrunkItemClass
	{
		ClsApi		= 1,
		ClsStrA		= 2,
		ClsStrW		= 4,
		ClsRef		= 0x80000000,
		ClsALL = -1,
	};

	static const char * filter_string [] = 
	{
		"padding",
		"<?xml version=",
		"<assembly",
		"assembly>",
		"processorarchitecture",
		"description>",
		"dependentassembly>",
		"assemblyidentity",
		"name=\"microsoft.windows.common-controls\"",
		"type=\"win32\"",
		"processorarchitecture",
		"requestedprivileges>",
		"security>",
		"trustinfo>",
		"dependency>",
		"<trustinfo",
		"requestedexecution",
		"windowssetting",
		"urn:schemas-microsoft-com:asm",
		"<application",
		"</application>",
		"\\varfileinfo\\translation",
		"\\stringfileinfo\\",
		"stringfileinfo",
		"<compatibility",
		"<supportedos id=",
		"</compatibility>",
		0,
	};

	static bool is_filter_mateched( const char * str )
	{
		for( size_t i = 0; filter_string[i]; ++ i )
		{
			if( tcs::str( str, filter_string[i] ) )
				return true;
		}
		return false;
	}


	struct TrunkItem
	{
		uint32_t		rva;
		uint32_t		tkls;
		uint32_t		foff;
		size_t			refed;
		std::string		text;
		bool_t			iscall;
		static bool _less_by_fpos( const TrunkItem & l, const TrunkItem & r ) 
		{
			return l.foff < r.foff;
		}
	};

#pragma pack(1)
	struct PSH_STUFF
	{
		uint16_t		machine;				// fileheader.machine
		uint32_t		charac;					// fileheader.characteristic(dll)
		uint16_t		num_rx_sect;			// 可读写的节个数
		uint16_t		num_rwx_sect;			// 可读可写的节个数

		uint16_t		magic;					// optheader.magic
		uint16_t		subsystem;				// from header

		uint16_t		oep_sec_idx;				// -1 : none resource
		uint32_t		oep_sec_attr;			// 

		uint16_t		res_sec_idx;			// -1 : none resource
		uint32_t		res_sec_attr;			// 

		uint16_t		iat_sec_idx;			// -1 : none import table,(可对付导入表在其他节重建的Patch)
		uint32_t		iat_sec_attr;			// 

		uint16_t		last_sec_idx;			// 最后一个节节号（相当于节个数）
		uint32_t		last_sec_attr;			// 最后一个节的属性

		uint16_t		dde_mask;				// with|without in 16 entries, ignore export-entry

		uint64_t		exp_sec_idx_set_crc;	// 所有导出函数入口所在的节号的集合的哈希（可应对某些导出函数Patch的情况）
		uint8_t			eps_1opcode_set[0x100];
		uint32_t		overlay_1m;				// 1m/2m/3m/4m

	} __attribute__((packed));
#pragma pack()

	class ThunkFetcher : string_dumper
	{
	protected:

		typedef std::map<uint32_t,TrunkItem>	TrunkItemMap;
		typedef std::vector<uint8_t>			ByteArray;

		TrunkItemMap			_item_map;
		refp<PeParser>			pep_;
		xm::u32wm::filter32_t	_filter32;
		uint64_t				_imgb;
		uint32_t				_imgs;
		std::string				_text;
		size_t					_done_mask;
		size_t					_last_text_gen;
		PSH_STUFF				psh_;

	public:

		ThunkFetcher() : _imgb(0), _imgs(0),_done_mask(0),_last_text_gen(0)
		{
			//Close();
		}

		size_t get_count()
		{
			return _item_map.size();
		}
		TrunkItem * get_next_item( uint32_t rva )
		{
			TrunkItemMap::iterator it = _item_map.lower_bound( rva );
			if( it == _item_map.end() ) return NULL;
			return &(it->second);
		}

		/*
		#define IMAGE_FILE_RELOCS_STRIPPED           0x0001  // Relocation info stripped from file.
		#define IMAGE_FILE_EXECUTABLE_IMAGE          0x0002  // File is executable  (i.e. no unresolved externel references).
		#define IMAGE_FILE_LINE_NUMS_STRIPPED        0x0004  // Line nunbers stripped from file.
		#define IMAGE_FILE_LOCAL_SYMS_STRIPPED       0x0008  // Local symbols stripped from file.
		#define IMAGE_FILE_AGGRESIVE_WS_TRIM         0x0010  // Agressively trim working set
		#define IMAGE_FILE_LARGE_ADDRESS_AWARE       0x0020  // App can handle >2gb addresses
		#define IMAGE_FILE_BYTES_REVERSED_LO         0x0080  // Bytes of machine word are reversed.
		#define IMAGE_FILE_32BIT_MACHINE             0x0100  // 32 bit word machine.
		#define IMAGE_FILE_DEBUG_STRIPPED            0x0200  // Debugging info stripped from file in .DBG file
		#define IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP   0x0400  // If Image is on removable media, copy and run from the swap file.
		#define IMAGE_FILE_NET_RUN_FROM_SWAP         0x0800  // If Image is on Net, copy and run from the swap file.
		#define IMAGE_FILE_SYSTEM                    0x1000  // System File.
		#define IMAGE_FILE_DLL                       0x2000  // File is a DLL.
		#define IMAGE_FILE_UP_SYSTEM_ONLY            0x4000  // File should only be run on a UP machine
		#define IMAGE_FILE_BYTES_REVERSED_HI         0x8000  // Bytes of machine word are reversed.
		*/

		inline void get_section_charac( uint16_t idx, uint32_t & charac )
		{
			if( idx == (uint16_t)-1 ) return ;
			WPE_IMAGE_SECTION_HEADER * pish = pep_->section_header(idx);
			if( !pish ) charac = 0;
			charac = pish->Characteristics & (WPE_IMAGE_SCN_MEM_EXECUTE|WPE_IMAGE_SCN_MEM_WRITE|WPE_IMAGE_SCN_CNT_CODE);
		}

		long get_psh( uint64_t * psh64 )
		{
			memset( &psh_, 0, sizeof(psh_) );

			psh_.machine = pep_->mfh.Machine;
			psh_.charac = pep_->mfh.Characteristics & (WPE_IMAGE_FILE_DLL);
			//////////////////////////////////////////////////////////////////////////
			// psh_.num_rwx_sect psh_.num_rx_sect
			size_t sn = pep_->num_of_sections_with_header();
			for( size_t i = 1; i < sn; ++ i )
			{
				WPE_IMAGE_SECTION_HEADER * header = pep_->section_header(i);
				if( !header ) break;
				bool_t is_code = header->Characteristics & (WPE_IMAGE_SCN_MEM_EXECUTE|WPE_IMAGE_SCN_CNT_CODE);
				bool_t is_writable = header->Characteristics & (WPE_IMAGE_SCN_MEM_WRITE);
				if( is_code && is_writable ) psh_.num_rwx_sect ++;
				if( is_code ) psh_.num_rx_sect ++;
			}
			//////////////////////////////////////////////////////////////////////////
			uint32_t iat_rva = 0, res_rva = 0, oep_rva = 0;
			//////////////////////////////////////////////////////////////////////////
			if( pep_->is_64bit() )
			{
				psh_.magic = pep_->moh64.Magic;
				psh_.subsystem = pep_->moh64.Subsystem;
				//////////////////////////////////////////////////////////////////////////
				size_t dde_num = MIN(16, pep_->moh64.NumberOfRvaAndSizes);
				for( size_t i = 0; i < dde_num; ++ i )
				{
					if( pep_->moh64.DataDirectory[i].VirtualAddress )
						psh_.dde_mask |= (1 << i);
				}
				//////////////////////////////////////////////////////////////////////////
				if( dde_num > WPE_IMAGE_DIRECTORY_ENTRY_IMPORT )
					iat_rva = pep_->moh64.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
				if( dde_num > WPE_IMAGE_DIRECTORY_ENTRY_RESOURCE )
					res_rva = pep_->moh64.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress;
				oep_rva = pep_->moh64.AddressOfEntryPoint;
			}
			else
			{
				psh_.magic = pep_->moh32.Magic;
				psh_.subsystem = pep_->moh32.Subsystem;
				//////////////////////////////////////////////////////////////////////////
				size_t dde_num = MIN(16, pep_->moh32.NumberOfRvaAndSizes);
				for( size_t i = 0; i < dde_num; ++ i )
				{
					if( pep_->moh32.DataDirectory[i].VirtualAddress )
						psh_.dde_mask |= (1 << i);
				}
				//////////////////////////////////////////////////////////////////////////
				if( dde_num > WPE_IMAGE_DIRECTORY_ENTRY_IMPORT )
					iat_rva = pep_->moh32.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
				if( dde_num > WPE_IMAGE_DIRECTORY_ENTRY_RESOURCE )
					res_rva = pep_->moh32.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress;
				oep_rva = pep_->moh32.AddressOfEntryPoint;
			}
			//////////////////////////////////////////////////////////////////////////
			if( oep_rva ) psh_.oep_sec_idx = pep_->get_section_index_by_rva( oep_rva ); else psh_.oep_sec_idx = -1;
			if( iat_rva ) psh_.iat_sec_idx = pep_->get_section_index_by_rva( iat_rva ); else psh_.iat_sec_idx = -1;
			if( res_rva ) psh_.res_sec_idx = pep_->get_section_index_by_rva( res_rva ); else psh_.res_sec_idx = -1;
			//////////////////////////////////////////////////////////////////////////
			get_section_charac( psh_.oep_sec_idx, psh_.oep_sec_attr );
			get_section_charac( psh_.iat_sec_idx, psh_.iat_sec_attr );
			get_section_charac( psh_.res_sec_idx, psh_.res_sec_attr );
			//////////////////////////////////////////////////////////////////////////
			if( oep_rva )
			{
				uint8_t op1 = 0;
				long lr = pep_->imgread_sc( oep_rva, &op1, 1 );
				if( lr < 0 ) return E_FAIL;
				psh_.eps_1opcode_set[op1] = 1;
			}
			//////////////////////////////////////////////////////////////////////////
			std::vector<uint8_t> exp_sects;
			exp_sects.resize( pep_->num_of_sections_with_header() );
			size_t n = pep_->get_num_of_exports();
			for( size_t e = 0; e < n; ++ e )
			{
				uint32_t exp_rva_ = 0;
				long lr = pep_->get_export_address( e, &exp_rva_ );
				if( lr < 0 ) return lr;
				uint8_t op1 = 0;
				lr = pep_->imgread_sc( exp_rva_, &op1, 1 );
				if( lr < 0 ) continue;
				//psh_.eps_1opcode_set[op1] = 1;
				//////////////////////////////////////////////////////////////////////////
				ssize_t si = pep_->get_section_index_by_rva( exp_rva_ );
				if( si < 0 || si >= exp_sects.size() )
					continue;
				exp_sects[si] = 1;
			}
			psh_.exp_sec_idx_set_crc = 'EXP@';
			psh_.exp_sec_idx_set_crc <<= 32;
			psh_.exp_sec_idx_set_crc += 'SECT';
			psh_.exp_sec_idx_set_crc = update_crc64( psh_.exp_sec_idx_set_crc, exp_sects.data(), exp_sects.size() );
			//////////////////////////////////////////////////////////////////////////

			*psh64 = 'PSH6';
			*psh64 <<= 32;
			*psh64 |= '4CRC';
			*psh64 = update_crc64( *psh64, (const uint8_t*)&psh_, sizeof(psh_) );
			//////////////////////////////////////////////////////////////////////////
			return S_OK;
		}


		long fetch( )
		{
			if( _done_mask & ClsALL ) return S_OK;

			size_t cimp = pep_->num_of_imports();
			//if( !cimp ) return S_OK;
			// add import
			for( size_t i = 0; i < cimp; ++ i )
			{
				PeParser::impapi_t * impp = pep_->get_impapi( i );
				uint32_t foff = 0;
				long lr = pep_->rva_to_raw( impp->iat, foff );
				if( lr < 0 ) continue;

				TrunkItem item;
				item.foff = foff;
				item.rva = impp->iat;
				item.iscall = false;
				item.refed = 0;
				item.text = impp->api;
				item.tkls = ClsApi;

				_item_map[item.rva] = item;
			}

			handle_sections( ClsStrA|ClsStrW );

			if( pep_->is_x86_64() )
			{
				_filter32.init( cimp + _item_map.size() );

				// add import
				for( size_t i = 0; i < cimp; ++ i )
				{
					PeParser::impapi_t * impp = pep_->get_impapi( i );
					_filter32.add( &(impp->iat) );
				}
				// add string
				TrunkItemMap::iterator it = _item_map.begin();
				for( ; it != _item_map.end(); ++ it )
				{
					TrunkItem& item = it->second;
					_filter32.add( &item.rva );
				}

				handle_sections( ClsApi );
			}

			_done_mask |= -1;

			return S_OK;
		}
		long close()
		{
			_item_map.clear();
			pep_.dispose();
			_filter32.clear();
			_imgb = 0;
			_imgs = 0;
			_text.clear();
			_done_mask = 0;
			_last_text_gen = 0;
			return S_OK;
		}

		long open( PeParser * pep )
		{
			close();

			if( !pep ) return E_INVALIDARG;

			_imgb = pep->get_image_base();
			_imgs = pep->get_image_size();
			pep_ = pep;

			return S_OK;
		}

		long get_as_text( std::string & text, bool with_info = false )
		{
			text.clear();
			TrunkItemMap::iterator it = _item_map.begin();
			for( ; it != _item_map.end(); )
			{
				TrunkItem& item = it->second;
				if( !(item.tkls & ClsRef) )
				{
					++ it ;
					continue;
				}
				else if( item.refed == it->first )
				{
					TrunkItemMap::iterator itdel = it ++;
					_item_map.erase( itdel );
				}
				else
				{
					_item_map.erase( item.refed );
					++ it ;
				}
			}

			it = _item_map.begin();
			for( ; it != _item_map.end(); ++ it )
			{
				TrunkItem& item = it->second;

				const char * kls = "RF";
				switch( item.tkls )
				{
				case ClsStrA|ClsRef: kls = "RA"; break;
				case ClsStrW|ClsRef: kls = "RW"; break;
				case ClsStrA: kls = " W"; break;
				case ClsStrW: kls = " W"; break;
				}
				if( with_info )
					tcs::xprintf( text, ":%08X|%s|%s|%s\n", item.foff, kls, item.iscall?"jmp":"ref", item.text.c_str() );
				else
					tcs::xprintf( text, "%s", item.text.c_str() );
			}

			return S_OK;
		}

	protected:

		bool on_data_ref_found_rva( uint32_t rva, uint32_t drva, bool_t is_call )
		{
			// it's api ref
			uint32_t raw = 0;
			long lr = pep_->rva_to_raw( rva, raw );
			if( lr < 0 ) return false;

			PeParser::impapi_t * impi = pep_->find_impapi( drva );
			if( impi )
			{
				_item_map.erase( drva );

				TrunkItem item;
				item.rva	= rva;
				item.tkls	= (ClsApi|ClsRef);
				item.foff	= raw;
				item.refed	= drva;
				item.text	= impi->api;
				item.iscall = is_call;
				_item_map[rva] = item;

				return true;
			}
			
			TrunkItemMap::iterator it = _item_map.find( drva );
			if( it == _item_map.end() ) 
				return true;

			TrunkItem & str = it->second;

			if( str.tkls != ClsStrA && str.tkls != ClsStrW )
				return true;
			
			TrunkItem	item;
			item.rva	= rva;
			item.tkls	= str.tkls|ClsRef;
			item.foff	= raw;
			item.refed	= str.rva;
			item.text	= str.text;
			item.iscall = false;
			_item_map[rva] = item;

			return true;
		}
		void on_data_ref_found_va( uint32_t rva, uint64_t dva, bool_t iscall = false )
		{
			if( dva < _imgb || dva >= (_imgb + _imgs ) ) return ;
			uint32_t iat = (uint32_t)(dva - _imgb);
			on_data_ref_found_rva( rva, iat, iscall );
		}

		virtual bool on_discover_string( size_t rva, char* str, size_t cc, size_t crc, bool utf16 )
		{
			//printf( "%08X : %s\n", rva, str );
			if( pep_->is_rva_of_import_name( (uint32_t)rva ) )
				return true;

			float si = 1;
			for( size_t i = 1; i < cc; ++ i )
			{
				if( str[i-1] == str[i] ) si ++;
			}
			if( si / (float)cc > 0.4 ) 
				return true;

			if( si == cc ) return true;

			uint32_t raw = 0;
			if( pep_->rva_to_raw( (uint32_t)rva, raw ) < 0 )
				return true;

			tcs::lwr( str );

			if( is_filter_mateched( str ) )
				return true;

			TrunkItem item;
			item.foff	= raw;
			item.rva	= rva;
			item.refed	= crc;
			item.text	= str;
			item.tkls	= utf16 ? ClsStrW : ClsStrA;
			item.iscall = false;

			_item_map[rva] = item;
			
			return true;
		}

		void handle_strdmp( uint32_t rvab, const uint8_t* bstart, const uint8_t* bend, size_t mask ) 
		{
			string_dumper::prepare( 4, 512 );
			string_dumper::dump_mbs( rvab, bstart, bend );
			string_dumper::dump_unis( rvab, bstart, bend );
			string_dumper::complete();
		}

		//
		//	搜索IAT，如果filter按照VA来构建，加速效果会更好，但是需要为32位和64位分别建立filter，所以目前先采用RVA构建filter
		//


		//
		//	64位代码，直接采用暴搜FF15/FF25和64位IAT
		//

		void handle_iat64( uint32_t foff, uint32_t rvab, const uint8_t* bstart, const uint8_t* bend ) 
		{
			size_t blen = bend - bstart;
			if( blen < sizeof(uint64_t) ) return ;

			bend -= sizeof(uint64_t);

			size_t skip = 1;
			uint64_t imge = _imgb + _imgs;

			for( const uint8_t* scan = bstart; scan < bend; scan += skip )
			{
				skip = 1;
				uint16_t op16 = *(le::U16*)scan;
				uint32_t nowrva = rvab + (uint32_t)(scan - bstart);

				if( op16 == 0x15FF || op16 == 0x25FF )
				{
					int32_t dist = *(le::I32*)(scan + sizeof(uint16_t));
					uint32_t _iat = nowrva + 6 + dist;
					if( !_filter32.test( &_iat ) )
					{
						if( on_data_ref_found_rva( nowrva + 2, _iat, true ) )
						{
							skip = 6;
							continue;
						}
					}
				}
				uint64_t iat = *(uint64_t*)scan;
				if( iat < _imgb ) continue;
				if( iat >= imge ) continue;
				iat -= _imgb;
				uint32_t _iat = (uint32_t)iat;
				if( !_filter32.test( &_iat ) )
				{
					if( on_data_ref_found_rva( nowrva, _iat, false ) )
					{
						skip += sizeof(uint64_t);
						continue;
					}
				}
			}
		}

		//
		//	32位代码，直接采用暴搜32位IAT
		//

		void handle_iat32( uint32_t foff, uint32_t rvab, const uint8_t* bstart, const uint8_t* bend ) 
		{
			size_t blen = bend - bstart;
			if( blen < sizeof(uint32_t) ) return ;

			bend -= sizeof(uint32_t);

			size_t skip = 1;
			uint64_t imge = _imgb + _imgs;
			
			for( const uint8_t* scan = bstart; scan < bend; scan += skip )
			{
				skip = 1;
				uint32_t iat = *(uint32_t*)scan;
				if( iat < _imgb ) continue;
				if( iat >= imge ) continue;
				iat -= (uint32_t)_imgb;
				skip = _filter32.test( &iat );
				if( skip ) continue;
				skip = 1;
				uint32_t off = (uint32_t)(scan - bstart);
				bool_t is_call = false;
				if( off > 2 )
				{
					uint16_t op16 = *(le::U16*)(scan-2);
					is_call = ( op16 == 0x15FF || op16 == 0x25FF );
				}
				uint32_t now_rva = rvab + off;
				//if( now_rva == 0x2EE2A ) __asm int 3;
				if( on_data_ref_found_rva( now_rva, iat, is_call ) )
					skip = sizeof(uint32_t);
			}
		}

		enum { MAX_CODE_SECTION_DATA_TO_HANDLE = 4 * 1024 * 1024 };
		enum { MAX_DATA_SECTION_DATA_TO_HANDLE = 32 * 1024 * 1024 };

		long handle_one_section( WPE_IMAGE_SECTION_HEADER * hdr, size_t mask )
		{
			// dump string
			size_t dumpcb = MIN( (uint32_t)hdr->SizeOfRawData, MAX_DATA_SECTION_DATA_TO_HANDLE );
			if( !dumpcb ) return S_OK;

			UTIL::sentry<UINT8*> buf( new UINT8[dumpcb] );
			if( !buf ) return E_OUTOFMEMORY;

			size_t eff = 0;
			long lr = pep_->get_fo()->readp( hdr->PointerToRawData, buf, dumpcb, &eff );
			if( eff != dumpcb ) return E_FAIL;

			cuchar_t * start = buf, * bend = buf + eff;

			if( pep_->is_x86_64() && (mask & ClsApi) )
			{
				if( pep_->is_64bit() )
					handle_iat64( hdr->PointerToRawData, hdr->VirtualAddress, start, bend );
				else
					handle_iat32( hdr->PointerToRawData, hdr->VirtualAddress, start, bend );
			}

			if( mask & (ClsStrA|ClsStrW) )
				handle_strdmp( hdr->VirtualAddress, start, bend, mask );

			return S_OK;
		}

		long handle_sections( size_t mask )
		{
			size_t sn = pep_->num_of_sections_with_header();

			for( DWORD i = 1; i < sn; ++ i )
			{
				WPE_IMAGE_SECTION_HEADER * hdr = pep_->section_header(i);
				if( !hdr ) break;

				if( !hdr->VirtualAddress ) continue;

				if( hdr->SizeOfRawData < 0x100 ) continue;

				handle_one_section( hdr, mask );
			}
			return S_OK;
		}


	
	};


};



#endif